{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "T5 on TPU",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyNhznFJcPoY54Wwqs6aEdRJ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "44ca82b3bbc5432eadc4f6fa3b81f483": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_9a58f575a463454aba3a52742abf00bf",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_6217a8ed0c634a5bbcbb4070d61d2ce3",
              "IPY_MODEL_9542805d664d451aa33a5bd7cc36d614"
            ]
          }
        },
        "9a58f575a463454aba3a52742abf00bf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "6217a8ed0c634a5bbcbb4070d61d2ce3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_f90e54ee9dee42a99f38fdf2186cea62",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 791656,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 791656,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_2085a86f64a04a1586fefa0fe1413406"
          }
        },
        "9542805d664d451aa33a5bd7cc36d614": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_8e0c79e36d8c4d1faab99cc565feeb17",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 792k/792k [00:06&lt;00:00, 124kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_413ae83fdf9d4573b2c2dc053db83aa2"
          }
        },
        "f90e54ee9dee42a99f38fdf2186cea62": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "2085a86f64a04a1586fefa0fe1413406": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "8e0c79e36d8c4d1faab99cc565feeb17": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "413ae83fdf9d4573b2c2dc053db83aa2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d7fa301d6f7d498bac466849f04321c5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_0f662bc0fa4444f386ec8975b4289f04",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_a2b08cf2bb834808b02d20fbee2f00ab",
              "IPY_MODEL_5efd16badcb64aa3a708975504feade1"
            ]
          }
        },
        "0f662bc0fa4444f386ec8975b4289f04": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a2b08cf2bb834808b02d20fbee2f00ab": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_a763d584af9147bcbe5430f95bf16925",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 4997,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 4997,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_329f27a50bee452b8f5ee719929fb33b"
          }
        },
        "5efd16badcb64aa3a708975504feade1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_0be03081b1324f79a99bdbe926c30e89",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 5.00k/5.00k [00:00&lt;00:00, 11.7kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_dca7ff44a73a4c9f872cb2596c3c729e"
          }
        },
        "a763d584af9147bcbe5430f95bf16925": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "329f27a50bee452b8f5ee719929fb33b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "0be03081b1324f79a99bdbe926c30e89": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "dca7ff44a73a4c9f872cb2596c3c729e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d530b763a0e9415f9b8b1f1787a66619": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_7b85b787b5d9490c994f9e406835813b",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_f86c72c0da864856b9ddb55568fe4a2b",
              "IPY_MODEL_a4e0f5b0b3cf4091a44bad74e5abb62e"
            ]
          }
        },
        "7b85b787b5d9490c994f9e406835813b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "f86c72c0da864856b9ddb55568fe4a2b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_3ecc3e6a31db4835987c22b81a85c23f",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 2240,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 2240,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_7c062b51acab4bb59eeb4269d5e6c2df"
          }
        },
        "a4e0f5b0b3cf4091a44bad74e5abb62e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_7fbf906e11d24bfdb90e3f0ce46250c1",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 2.24k/2.24k [00:02&lt;00:00, 1.02kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_43deb763aee24506a82dd795cd126f49"
          }
        },
        "3ecc3e6a31db4835987c22b81a85c23f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "7c062b51acab4bb59eeb4269d5e6c2df": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "7fbf906e11d24bfdb90e3f0ce46250c1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "43deb763aee24506a82dd795cd126f49": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "6af39625f78e4e5989739eff7e8849f6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_d97ac4c36e7944fdb989cf7ee1851350",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_2539f0244f0e468a993609fd562ff1f5",
              "IPY_MODEL_17025b38dedb45a39a95ffe7b6117d87"
            ]
          }
        },
        "d97ac4c36e7944fdb989cf7ee1851350": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2539f0244f0e468a993609fd562ff1f5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_b6dde636eb51433fb5be7af2cedc2774",
            "_dom_classes": [],
            "description": "Downloading: ",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 8116577,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 8116577,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_ff877f980dd34d2d9f3a4d3efb50af0c"
          }
        },
        "17025b38dedb45a39a95ffe7b6117d87": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_cde27b62215448be92d8c8fa2292438a",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 30.3M/? [00:00&lt;00:00, 41.3MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_5b66084b700145d5903aa91bcd2920f7"
          }
        },
        "b6dde636eb51433fb5be7af2cedc2774": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "ff877f980dd34d2d9f3a4d3efb50af0c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "cde27b62215448be92d8c8fa2292438a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "5b66084b700145d5903aa91bcd2920f7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "741a693169b54851a4cb47369dd9bd1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_eb0f02133b9140be92288fd6a8415deb",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_03a75d6d64524d47a3df6457e1901cbd",
              "IPY_MODEL_b4a2213a0a0b421d895269733344cc9d"
            ]
          }
        },
        "eb0f02133b9140be92288fd6a8415deb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "03a75d6d64524d47a3df6457e1901cbd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_64e947e3d9e440cf9f3880cb2325e57a",
            "_dom_classes": [],
            "description": "Downloading: ",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1054280,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1054280,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_108bbc9991614357aae4469fc19b4859"
          }
        },
        "b4a2213a0a0b421d895269733344cc9d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_6792931cb55944deb260f07d6bf87931",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 4.85M/? [00:00&lt;00:00, 15.2MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3cb18f41264740c9933f159a9820ad77"
          }
        },
        "64e947e3d9e440cf9f3880cb2325e57a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "108bbc9991614357aae4469fc19b4859": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "6792931cb55944deb260f07d6bf87931": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "3cb18f41264740c9933f159a9820ad77": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "004faf88da144ad2bd3e50b5b9766627": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_f02e5f34ac414081ab2882b2ab7ff64d",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_41c3aac64e424e1db59b723a11bb0fb0",
              "IPY_MODEL_44540a280fb141979bdeff7f93d0cde2"
            ]
          }
        },
        "f02e5f34ac414081ab2882b2ab7ff64d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "41c3aac64e424e1db59b723a11bb0fb0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_f9c604fa9a8642eaa218824b8ab651fe",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 1,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_34e5dacca6fc4497abda0a2ad37996e0"
          }
        },
        "44540a280fb141979bdeff7f93d0cde2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_8bd88e5bf56e45bbbc1d339b00eff57c",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 87599/0 [00:03&lt;00:00, 5389.75 examples/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_93b33f1c33de44b4a1e57fb4133343c9"
          }
        },
        "f9c604fa9a8642eaa218824b8ab651fe": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "34e5dacca6fc4497abda0a2ad37996e0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "8bd88e5bf56e45bbbc1d339b00eff57c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "93b33f1c33de44b4a1e57fb4133343c9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3ec93bdea33b40ba82e24b4a0e605bc0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_121c3099d6e04079a4f6bf5707e0c60d",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_4f17790e2f594811bbfe55c9d883cf41",
              "IPY_MODEL_602e1a740e8f41fb800dbb3af0ea04f9"
            ]
          }
        },
        "121c3099d6e04079a4f6bf5707e0c60d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4f17790e2f594811bbfe55c9d883cf41": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_c1628c0657a44cf0873f2d64c9fb549c",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 1,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_99f6bf35ae7f4e66b4d339fe93768c7c"
          }
        },
        "602e1a740e8f41fb800dbb3af0ea04f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d73cf402af074a5481bf900e6abe2c63",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 10570/0 [00:00&lt;00:00, 18381.93 examples/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_0b2c415901d8468fb3f9acb63dd2f4e3"
          }
        },
        "c1628c0657a44cf0873f2d64c9fb549c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "99f6bf35ae7f4e66b4d339fe93768c7c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d73cf402af074a5481bf900e6abe2c63": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "0b2c415901d8468fb3f9acb63dd2f4e3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "1d197b0946534a91b2eb4595baade174": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_55c18e9169ab44c981de4cb2ebaea578",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_21065bd4f9fa4d97a20a25c9a35b7737",
              "IPY_MODEL_88d28cee2f02428aa3c4c22a873d6385"
            ]
          }
        },
        "55c18e9169ab44c981de4cb2ebaea578": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "21065bd4f9fa4d97a20a25c9a35b7737": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_12d2d5b9a3f54b25aa8fef6439625db6",
            "_dom_classes": [],
            "description": "Epoch: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 4,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 4,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_32f70451efa6473a8a00807dd22a3248"
          }
        },
        "88d28cee2f02428aa3c4c22a873d6385": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_17c358e37fc346a9ac8ce475599e495a",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 4/4 [20:24&lt;00:00, 306.06s/it]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_947afd4c522443b38bc223114178c7ee"
          }
        },
        "12d2d5b9a3f54b25aa8fef6439625db6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "32f70451efa6473a8a00807dd22a3248": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "17c358e37fc346a9ac8ce475599e495a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "947afd4c522443b38bc223114178c7ee": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a558cfceca8b43568dfad2267a534a35": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_7a5270ddb66447af9fba298e830fe8a4",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_d961f9fa40a04ceda3e1962eef4d33cc",
              "IPY_MODEL_d9363dcd670643e6ad2d0b65c4238b98"
            ]
          }
        },
        "7a5270ddb66447af9fba298e830fe8a4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d961f9fa40a04ceda3e1962eef4d33cc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_b25b3d9394be4ac786678181c800e926",
            "_dom_classes": [],
            "description": "Iteration: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1369,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1369,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3fed14bdec8b4162887d77dea7a9f289"
          }
        },
        "d9363dcd670643e6ad2d0b65c4238b98": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_5051978cf84f4f8191c1e3ad4661776e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1369/1369 [20:23&lt;00:00,  1.12it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_f67b49a0d48f4714a124bc972e9afa0f"
          }
        },
        "b25b3d9394be4ac786678181c800e926": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "3fed14bdec8b4162887d77dea7a9f289": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "5051978cf84f4f8191c1e3ad4661776e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "f67b49a0d48f4714a124bc972e9afa0f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2c4d7bdf7f6c48f0825ef87a12399f39": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_391fc0937c3b48f8b4c3b33a1c3849a0",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_11c6ff7dfb19436c920c5114734c0ec9",
              "IPY_MODEL_3b3722a55e3d43bc93c29a505895a25c"
            ]
          }
        },
        "391fc0937c3b48f8b4c3b33a1c3849a0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "11c6ff7dfb19436c920c5114734c0ec9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_f46743ac2da84d6ca77180d8686ead6b",
            "_dom_classes": [],
            "description": "Iteration:   0%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 1369,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_1b0e2b414d6843ae9c918db417ffff4d"
          }
        },
        "3b3722a55e3d43bc93c29a505895a25c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_e7559a6907f7479db7764f80c4501993",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/1369 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_5f89a141b34d478e861ea6bd4a17dcd2"
          }
        },
        "f46743ac2da84d6ca77180d8686ead6b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "1b0e2b414d6843ae9c918db417ffff4d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "e7559a6907f7479db7764f80c4501993": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "5f89a141b34d478e861ea6bd4a17dcd2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a9a7aeae6fa74b67a349a69e845f8898": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_0a9d6c739f984b91a0e17724707eb7e1",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_1ed37d3716374294b072524937625965",
              "IPY_MODEL_4b44de48aeeb4b8dbc9388fc7f7da227"
            ]
          }
        },
        "0a9d6c739f984b91a0e17724707eb7e1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "1ed37d3716374294b072524937625965": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_051052ef94d447a4b0233225d8f5ae81",
            "_dom_classes": [],
            "description": "Iteration:   0%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 1369,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3e43d0668e6f461c9cc461aa735a4437"
          }
        },
        "4b44de48aeeb4b8dbc9388fc7f7da227": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_43e265f386ca4dfd88e226a87f32775c",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/1369 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_6568a6e42f734c29ad5c26f2397a5a69"
          }
        },
        "051052ef94d447a4b0233225d8f5ae81": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "3e43d0668e6f461c9cc461aa735a4437": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "43e265f386ca4dfd88e226a87f32775c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "6568a6e42f734c29ad5c26f2397a5a69": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "c227333162114432a38b1480ec837701": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_71a3e5b846a74f409dfd029ea3bf5d61",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_8a27b38cd17649b1bc720a33050dc520",
              "IPY_MODEL_16497f2b4de443bb9bb14ccaf0c94803"
            ]
          }
        },
        "71a3e5b846a74f409dfd029ea3bf5d61": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "8a27b38cd17649b1bc720a33050dc520": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_f93601a0c6d84d8b9324fa2380e14d25",
            "_dom_classes": [],
            "description": "Iteration:   0%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 1369,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_f30f39adee694c55bcbaa850cf031b69"
          }
        },
        "16497f2b4de443bb9bb14ccaf0c94803": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_ea4e4964058945f287eb2a6271b5bae7",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/1369 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_1e74937c1f304d15ad49230b2428590d"
          }
        },
        "f93601a0c6d84d8b9324fa2380e14d25": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "f30f39adee694c55bcbaa850cf031b69": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ea4e4964058945f287eb2a6271b5bae7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "1e74937c1f304d15ad49230b2428590d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "df609e7f302e48949a49fae1b0deeb54": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_be907685a2b949d0a088a907a2cb2a90",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_743bba68cbe943d78ce63a1efd70d929",
              "IPY_MODEL_f3fffc31462b49e193b3c24d600b4ddb"
            ]
          }
        },
        "be907685a2b949d0a088a907a2cb2a90": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "743bba68cbe943d78ce63a1efd70d929": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_87fa58101b24400da90e4b0f87c7cb1a",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 331,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 331,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3c39d97054ba47a3ad3d9bb9a882756e"
          }
        },
        "f3fffc31462b49e193b3c24d600b4ddb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_20855b3cb56944cfb0d6aa13bd06429a",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 331/331 [1:31:09&lt;00:00, 16.52s/it]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_779d80cc14104171ab40ff53170e4807"
          }
        },
        "87fa58101b24400da90e4b0f87c7cb1a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "3c39d97054ba47a3ad3d9bb9a882756e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "20855b3cb56944cfb0d6aa13bd06429a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "779d80cc14104171ab40ff53170e4807": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_6q2San5Xh5N",
        "colab_type": "text"
      },
      "source": [
        "# T5 on TPU 💥🚀"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "knpacarPX2AN",
        "colab_type": "text"
      },
      "source": [
        "In this notebook we will see how to train T5 model on TPU with Huggingface's awesome new [trainer](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py). We will train T5 base model on SQUAD dataset for QA task. We will use the recently released amazing [nlp](https://github.com/huggingface/nlp) package to load and process the dataset in just few lines.\n",
        "\n",
        "First make sure you are connected to the high RAM instance. This will not work on 12 GB colab instance."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QLGiFCDqvuil",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Crash on purpose to get more ram :\n",
        "import torch\n",
        "torch.tensor([10.]*10000000000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40CJSrN9ZiIP",
        "colab_type": "text"
      },
      "source": [
        "Let's install [PyTorch/XLA](https://github.com/pytorch/xla) which enables PyTorch on TPU. Make sure you install the nightly version, as the trainer breaks on other versions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "coOmS2s_xDBy",
        "colab_type": "code",
        "outputId": "6ed0947e-4061-4ba7-9eca-72e7ed935bd6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 907
        }
      },
      "source": [
        "VERSION = \"nightly\"  #@param [\"1.5\" , \"20200325\", \"nightly\"]\n",
        "!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
        "!python pytorch-xla-env-setup.py --version $VERSION"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "\r  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\r100  4264  100  4264    0     0  60914      0 --:--:-- --:--:-- --:--:-- 60914\n",
            "Updating TPU and VM. This may take around 2 minutes.\n",
            "Updating TPU runtime to pytorch-nightly ...\n",
            "Uninstalling torch-1.5.0+cu101:\n",
            "Done updating TPU runtime: <Response [200]>\n",
            "  Successfully uninstalled torch-1.5.0+cu101\n",
            "Uninstalling torchvision-0.6.0+cu101:\n",
            "  Successfully uninstalled torchvision-0.6.0+cu101\n",
            "Copying gs://tpu-pytorch/wheels/torch-nightly-cp36-cp36m-linux_x86_64.whl...\n",
            "- [1 files][ 91.1 MiB/ 91.1 MiB]                                                \n",
            "Operation completed over 1 objects/91.1 MiB.                                     \n",
            "Copying gs://tpu-pytorch/wheels/torch_xla-nightly-cp36-cp36m-linux_x86_64.whl...\n",
            "- [1 files][119.8 MiB/119.8 MiB]                                                \n",
            "Operation completed over 1 objects/119.8 MiB.                                    \n",
            "Copying gs://tpu-pytorch/wheels/torchvision-nightly-cp36-cp36m-linux_x86_64.whl...\n",
            "/ [1 files][  2.3 MiB/  2.3 MiB]                                                \n",
            "Operation completed over 1 objects/2.3 MiB.                                      \n",
            "Processing ./torch-nightly-cp36-cp36m-linux_x86_64.whl\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==nightly) (1.18.4)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch==nightly) (0.16.0)\n",
            "\u001b[31mERROR: fastai 1.0.61 requires torchvision, which is not installed.\u001b[0m\n",
            "Installing collected packages: torch\n",
            "Successfully installed torch-1.6.0a0+83df3be\n",
            "Processing ./torch_xla-nightly-cp36-cp36m-linux_x86_64.whl\n",
            "Installing collected packages: torch-xla\n",
            "Successfully installed torch-xla-1.6+2191422\n",
            "Processing ./torchvision-nightly-cp36-cp36m-linux_x86_64.whl\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly) (7.0.0)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly) (1.6.0a0+83df3be)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly) (1.18.4)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch->torchvision==nightly) (0.16.0)\n",
            "Installing collected packages: torchvision\n",
            "Successfully installed torchvision-0.7.0a0+348dd5a\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following NEW packages will be installed:\n",
            "  libomp5\n",
            "0 upgraded, 1 newly installed, 0 to remove and 31 not upgraded.\n",
            "Need to get 234 kB of archives.\n",
            "After this operation, 774 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic/universe amd64 libomp5 amd64 5.0.1-1 [234 kB]\n",
            "Fetched 234 kB in 1s (371 kB/s)\n",
            "Selecting previously unselected package libomp5:amd64.\n",
            "(Reading database ... 144433 files and directories currently installed.)\n",
            "Preparing to unpack .../libomp5_5.0.1-1_amd64.deb ...\n",
            "Unpacking libomp5:amd64 (5.0.1-1) ...\n",
            "Setting up libomp5:amd64 (5.0.1-1) ...\n",
            "Processing triggers for libc-bin (2.27-3ubuntu1) ...\n",
            "/sbin/ldconfig.real: /usr/local/lib/python3.6/dist-packages/ideep4py/lib/libmkldnn.so.0 is not a symbolic link\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6-p80vyFZ-S8",
        "colab_type": "text"
      },
      "source": [
        "Install transformers and the nlp package. Restart colab after this"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ptPupnLsfkMH",
        "colab_type": "code",
        "outputId": "94ed3f2e-a762-4969-cc6f-34382fc779ec",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "!git clone https://github.com/huggingface/transformers.git\n",
        "!pip install ./transformers\n",
        "!pip install -U nlp"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'transformers'...\n",
            "remote: Enumerating objects: 94, done.\u001b[K\n",
            "remote: Counting objects:   1% (1/94)\u001b[K\rremote: Counting objects:   2% (2/94)\u001b[K\rremote: Counting objects:   3% (3/94)\u001b[K\rremote: Counting objects:   4% (4/94)\u001b[K\rremote: Counting objects:   5% (5/94)\u001b[K\rremote: Counting objects:   6% (6/94)\u001b[K\rremote: Counting objects:   7% (7/94)\u001b[K\rremote: Counting objects:   8% (8/94)\u001b[K\rremote: Counting objects:   9% (9/94)\u001b[K\rremote: Counting objects:  10% (10/94)\u001b[K\rremote: Counting objects:  11% (11/94)\u001b[K\rremote: Counting objects:  12% (12/94)\u001b[K\rremote: Counting objects:  13% (13/94)\u001b[K\rremote: Counting objects:  14% (14/94)\u001b[K\rremote: Counting objects:  15% (15/94)\u001b[K\rremote: Counting objects:  17% (16/94)\u001b[K\rremote: Counting objects:  18% (17/94)\u001b[K\rremote: Counting objects:  19% (18/94)\u001b[K\rremote: Counting objects:  20% (19/94)\u001b[K\rremote: Counting objects:  21% (20/94)\u001b[K\rremote: Counting objects:  22% (21/94)\u001b[K\rremote: Counting objects:  23% (22/94)\u001b[K\rremote: Counting objects:  24% (23/94)\u001b[K\rremote: Counting objects:  25% (24/94)\u001b[K\rremote: Counting objects:  26% (25/94)\u001b[K\rremote: Counting objects:  27% (26/94)\u001b[K\rremote: Counting objects:  28% (27/94)\u001b[K\rremote: Counting objects:  29% (28/94)\u001b[K\rremote: Counting objects:  30% (29/94)\u001b[K\rremote: Counting objects:  31% (30/94)\u001b[K\rremote: Counting objects:  32% (31/94)\u001b[K\rremote: Counting objects:  34% (32/94)\u001b[K\rremote: Counting objects:  35% (33/94)\u001b[K\rremote: Counting objects:  36% (34/94)\u001b[K\rremote: Counting objects:  37% (35/94)\u001b[K\rremote: Counting objects:  38% (36/94)\u001b[K\rremote: Counting objects:  39% (37/94)\u001b[K\rremote: Counting objects:  40% (38/94)\u001b[K\rremote: Counting objects:  41% (39/94)\u001b[K\rremote: Counting objects:  42% (40/94)\u001b[K\rremote: Counting objects:  43% (41/94)\u001b[K\rremote: Counting objects:  44% (42/94)\u001b[K\rremote: Counting objects:  45% (43/94)\u001b[K\rremote: Counting objects:  46% (44/94)\u001b[K\rremote: Counting objects:  47% (45/94)\u001b[K\rremote: Counting objects:  48% (46/94)\u001b[K\rremote: Counting objects:  50% (47/94)\u001b[K\rremote: Counting objects:  51% (48/94)\u001b[K\rremote: Counting objects:  52% (49/94)\u001b[K\rremote: Counting objects:  53% (50/94)\u001b[K\rremote: Counting objects:  54% (51/94)\u001b[K\rremote: Counting objects:  55% (52/94)\u001b[K\rremote: Counting objects:  56% (53/94)\u001b[K\rremote: Counting objects:  57% (54/94)\u001b[K\rremote: Counting objects:  58% (55/94)\u001b[K\rremote: Counting objects:  59% (56/94)\u001b[K\rremote: Counting objects:  60% (57/94)\u001b[K\rremote: Counting objects:  61% (58/94)\u001b[K\rremote: Counting objects:  62% (59/94)\u001b[K\rremote: Counting objects:  63% (60/94)\u001b[K\rremote: Counting objects:  64% (61/94)\u001b[K\rremote: Counting objects:  65% (62/94)\u001b[K\rremote: Counting objects:  67% (63/94)\u001b[K\rremote: Counting objects:  68% (64/94)\u001b[K\rremote: Counting objects:  69% (65/94)\u001b[K\rremote: Counting objects:  70% (66/94)\u001b[K\rremote: Counting objects:  71% (67/94)\u001b[K\rremote: Counting objects:  72% (68/94)\u001b[K\rremote: Counting objects:  73% (69/94)\u001b[K\rremote: Counting objects:  74% (70/94)\u001b[K\rremote: Counting objects:  75% (71/94)\u001b[K\rremote: Counting objects:  76% (72/94)\u001b[K\rremote: Counting objects:  77% (73/94)\u001b[K\rremote: Counting objects:  78% (74/94)\u001b[K\rremote: Counting objects:  79% (75/94)\u001b[K\rremote: Counting objects:  80% (76/94)\u001b[K\rremote: Counting objects:  81% (77/94)\u001b[K\rremote: Counting objects:  82% (78/94)\u001b[K\rremote: Counting objects:  84% (79/94)\u001b[K\rremote: Counting objects:  85% (80/94)\u001b[K\rremote: Counting objects:  86% (81/94)\u001b[K\rremote: Counting objects:  87% (82/94)\u001b[K\rremote: Counting objects:  88% (83/94)\u001b[K\rremote: Counting objects:  89% (84/94)\u001b[K\rremote: Counting objects:  90% (85/94)\u001b[K\rremote: Counting objects:  91% (86/94)\u001b[K\rremote: Counting objects:  92% (87/94)\u001b[K\rremote: Counting objects:  93% (88/94)\u001b[K\rremote: Counting objects:  94% (89/94)\u001b[K\rremote: Counting objects:  95% (90/94)\u001b[K\rremote: Counting objects:  96% (91/94)\u001b[K\rremote: Counting objects:  97% (92/94)\u001b[K\rremote: Counting objects:  98% (93/94)\u001b[K\rremote: Counting objects: 100% (94/94)\u001b[K\rremote: Counting objects: 100% (94/94), done.\u001b[K\n",
            "remote: Compressing objects: 100% (46/46), done.\u001b[K\n",
            "remote: Total 26447 (delta 52), reused 66 (delta 32), pack-reused 26353\u001b[K\n",
            "Receiving objects: 100% (26447/26447), 15.95 MiB | 29.37 MiB/s, done.\n",
            "Resolving deltas: 100% (18423/18423), done.\n",
            "Processing ./transformers\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers==2.9.1) (1.18.4)\n",
            "Collecting tokenizers==0.7.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/14/e5/a26eb4716523808bb0a799fcfdceb6ebf77a18169d9591b2f46a9adb87d9/tokenizers-0.7.0-cp36-cp36m-manylinux1_x86_64.whl (3.8MB)\n",
            "\u001b[K     |████████████████████████████████| 3.8MB 3.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.6/dist-packages (from transformers==2.9.1) (3.0.12)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers==2.9.1) (2.23.0)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.6/dist-packages (from transformers==2.9.1) (4.41.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.6/dist-packages (from transformers==2.9.1) (2019.12.20)\n",
            "Collecting sentencepiece\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/3b/88/49e772d686088e1278766ad68a463513642a2a877487decbd691dec02955/sentencepiece-0.1.90-cp36-cp36m-manylinux1_x86_64.whl (1.1MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1MB 32.5MB/s \n",
            "\u001b[?25hCollecting sacremoses\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/7d/34/09d19aff26edcc8eb2a01bed8e98f13a1537005d31e95233fd48216eed10/sacremoses-0.0.43.tar.gz (883kB)\n",
            "\u001b[K     |████████████████████████████████| 890kB 46.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: dataclasses in /usr/local/lib/python3.6/dist-packages (from transformers==2.9.1) (0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==2.9.1) (2020.4.5.1)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==2.9.1) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==2.9.1) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers==2.9.1) (2.9)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers==2.9.1) (1.12.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers==2.9.1) (7.1.2)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers==2.9.1) (0.14.1)\n",
            "Building wheels for collected packages: transformers, sacremoses\n",
            "  Building wheel for transformers (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for transformers: filename=transformers-2.9.1-cp36-none-any.whl size=637722 sha256=d72ddcc31cb33178d6e0190c642ac0d5fe63801c03375fcf7c85b379c12a2938\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-417x3jto/wheels/23/19/dd/2561a4e47240cf6b307729d58e56f8077dd0c698f5992216cf\n",
            "  Building wheel for sacremoses (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sacremoses: filename=sacremoses-0.0.43-cp36-none-any.whl size=893260 sha256=7d348f2205ae5ec3039f13604ba844a8aa8083f6b9a96382cb0b3ee487157340\n",
            "  Stored in directory: /root/.cache/pip/wheels/29/3c/fd/7ce5c3f0666dab31a50123635e6fb5e19ceb42ce38d4e58f45\n",
            "Successfully built transformers sacremoses\n",
            "Installing collected packages: tokenizers, sentencepiece, sacremoses, transformers\n",
            "Successfully installed sacremoses-0.0.43 sentencepiece-0.1.90 tokenizers-0.7.0 transformers-2.9.1\n",
            "Collecting nlp\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/21/17/0a408ac3403c71c978e7906146c69ed00b18712a4548e34d0ffb567c34cc/nlp-0.1.0-py3-none-any.whl (87kB)\n",
            "\u001b[K     |████████████████████████████████| 92kB 2.9MB/s \n",
            "\u001b[?25hRequirement already satisfied, skipping upgrade: numpy in /usr/local/lib/python3.6/dist-packages (from nlp) (1.18.4)\n",
            "Requirement already satisfied, skipping upgrade: tqdm>=4.27 in /usr/local/lib/python3.6/dist-packages (from nlp) (4.41.1)\n",
            "Collecting pyarrow>=0.16.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/8c/d8/ff07b4cf88362ae4bebe5b7ffc357e2e6c773cdf2db8fec1bf9e973012dd/pyarrow-0.17.0-cp36-cp36m-manylinux2014_x86_64.whl (63.8MB)\n",
            "\u001b[K     |████████████████████████████████| 63.8MB 69kB/s \n",
            "\u001b[?25hRequirement already satisfied, skipping upgrade: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from nlp) (0.7)\n",
            "Requirement already satisfied, skipping upgrade: filelock in /usr/local/lib/python3.6/dist-packages (from nlp) (3.0.12)\n",
            "Requirement already satisfied, skipping upgrade: requests>=2.19.0 in /usr/local/lib/python3.6/dist-packages (from nlp) (2.23.0)\n",
            "Requirement already satisfied, skipping upgrade: dill in /usr/local/lib/python3.6/dist-packages (from nlp) (0.3.1.1)\n",
            "Requirement already satisfied, skipping upgrade: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->nlp) (1.24.3)\n",
            "Requirement already satisfied, skipping upgrade: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->nlp) (2.9)\n",
            "Requirement already satisfied, skipping upgrade: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->nlp) (3.0.4)\n",
            "Requirement already satisfied, skipping upgrade: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->nlp) (2020.4.5.1)\n",
            "Installing collected packages: pyarrow, nlp\n",
            "  Found existing installation: pyarrow 0.14.1\n",
            "    Uninstalling pyarrow-0.14.1:\n",
            "      Successfully uninstalled pyarrow-0.14.1\n",
            "Successfully installed nlp-0.1.0 pyarrow-0.17.0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "pyarrow"
                ]
              }
            }
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "--2020-05-16 16:02:40--  https://raw.githubusercontent.com/huggingface/transformers/2d184cb553ee20943b03b253f44300e466357871/examples/xla_spawn.py\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 151.101.0.133, 151.101.64.133, 151.101.128.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|151.101.0.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1913 (1.9K) [text/plain]\n",
            "Saving to: ‘xla_spawn.py’\n",
            "\n",
            "\rxla_spawn.py          0%[                    ]       0  --.-KB/s               \rxla_spawn.py        100%[===================>]   1.87K  --.-KB/s    in 0s      \n",
            "\n",
            "2020-05-16 16:02:41 (40.3 MB/s) - ‘xla_spawn.py’ saved [1913/1913]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zFWlfEJllAcw",
        "colab_type": "text"
      },
      "source": [
        "## Load and process data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NVOz2QUtaKQb",
        "colab_type": "text"
      },
      "source": [
        "Let's load and process the dataset using the nlp library. We will process the examples in follwoing way to cast QA task in text-to-text setting\n",
        "\n",
        "**input**\n",
        "question: question_text  context: context \n",
        "\n",
        "**target**\n",
        "answer_text"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CaRw0ke1e1sF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import nlp\n",
        "from transformers import T5Tokenizer"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NaGYDvKUe8VS",
        "colab_type": "code",
        "outputId": "a8cba6aa-c83d-4efb-d5ee-9bc6831509af",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "44ca82b3bbc5432eadc4f6fa3b81f483",
            "9a58f575a463454aba3a52742abf00bf",
            "6217a8ed0c634a5bbcbb4070d61d2ce3",
            "9542805d664d451aa33a5bd7cc36d614",
            "f90e54ee9dee42a99f38fdf2186cea62",
            "2085a86f64a04a1586fefa0fe1413406",
            "8e0c79e36d8c4d1faab99cc565feeb17",
            "413ae83fdf9d4573b2c2dc053db83aa2"
          ]
        }
      },
      "source": [
        "tokenizer = T5Tokenizer.from_pretrained('t5-base')"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "44ca82b3bbc5432eadc4f6fa3b81f483",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=791656.0, style=ProgressStyle(descripti…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-gJOEe0Ye0di",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# process the examples in input and target text format and the eos token at the end \n",
        "def add_eos_to_examples(example):\n",
        "    example['input_text'] = 'question: %s  context: %s </s>' % (example['question'], example['context'])\n",
        "    example['target_text'] = '%s </s>' % example['answers']['text'][0]\n",
        "    return example\n",
        "\n",
        "# tokenize the examples\n",
        "def convert_to_features(example_batch):\n",
        "    input_encodings = tokenizer.batch_encode_plus(example_batch['input_text'], pad_to_max_length=True, max_length=512)\n",
        "    target_encodings = tokenizer.batch_encode_plus(example_batch['target_text'], pad_to_max_length=True, max_length=16)\n",
        "\n",
        "    encodings = {\n",
        "        'input_ids': input_encodings['input_ids'], \n",
        "        'attention_mask': input_encodings['attention_mask'],\n",
        "        'target_ids': target_encodings['input_ids'],\n",
        "        'target_attention_mask': target_encodings['attention_mask']\n",
        "    }\n",
        "\n",
        "    return encodings"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ZWE4addfSmi",
        "colab_type": "code",
        "outputId": "aee10e16-998e-45cd-feaa-e05d1f42283b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 333,
          "referenced_widgets": [
            "d7fa301d6f7d498bac466849f04321c5",
            "0f662bc0fa4444f386ec8975b4289f04",
            "a2b08cf2bb834808b02d20fbee2f00ab",
            "5efd16badcb64aa3a708975504feade1",
            "a763d584af9147bcbe5430f95bf16925",
            "329f27a50bee452b8f5ee719929fb33b",
            "0be03081b1324f79a99bdbe926c30e89",
            "dca7ff44a73a4c9f872cb2596c3c729e",
            "d530b763a0e9415f9b8b1f1787a66619",
            "7b85b787b5d9490c994f9e406835813b",
            "f86c72c0da864856b9ddb55568fe4a2b",
            "a4e0f5b0b3cf4091a44bad74e5abb62e",
            "3ecc3e6a31db4835987c22b81a85c23f",
            "7c062b51acab4bb59eeb4269d5e6c2df",
            "7fbf906e11d24bfdb90e3f0ce46250c1",
            "43deb763aee24506a82dd795cd126f49",
            "6af39625f78e4e5989739eff7e8849f6",
            "d97ac4c36e7944fdb989cf7ee1851350",
            "2539f0244f0e468a993609fd562ff1f5",
            "17025b38dedb45a39a95ffe7b6117d87",
            "b6dde636eb51433fb5be7af2cedc2774",
            "ff877f980dd34d2d9f3a4d3efb50af0c",
            "cde27b62215448be92d8c8fa2292438a",
            "5b66084b700145d5903aa91bcd2920f7",
            "741a693169b54851a4cb47369dd9bd1e",
            "eb0f02133b9140be92288fd6a8415deb",
            "03a75d6d64524d47a3df6457e1901cbd",
            "b4a2213a0a0b421d895269733344cc9d",
            "64e947e3d9e440cf9f3880cb2325e57a",
            "108bbc9991614357aae4469fc19b4859",
            "6792931cb55944deb260f07d6bf87931",
            "3cb18f41264740c9933f159a9820ad77",
            "004faf88da144ad2bd3e50b5b9766627",
            "f02e5f34ac414081ab2882b2ab7ff64d",
            "41c3aac64e424e1db59b723a11bb0fb0",
            "44540a280fb141979bdeff7f93d0cde2",
            "f9c604fa9a8642eaa218824b8ab651fe",
            "34e5dacca6fc4497abda0a2ad37996e0",
            "8bd88e5bf56e45bbbc1d339b00eff57c",
            "93b33f1c33de44b4a1e57fb4133343c9",
            "3ec93bdea33b40ba82e24b4a0e605bc0",
            "121c3099d6e04079a4f6bf5707e0c60d",
            "4f17790e2f594811bbfe55c9d883cf41",
            "602e1a740e8f41fb800dbb3af0ea04f9",
            "c1628c0657a44cf0873f2d64c9fb549c",
            "99f6bf35ae7f4e66b4d339fe93768c7c",
            "d73cf402af074a5481bf900e6abe2c63",
            "0b2c415901d8468fb3f9acb63dd2f4e3"
          ]
        }
      },
      "source": [
        "# load train and validation split of squad\n",
        "train_dataset  = nlp.load_dataset('squad', split=nlp.Split.TRAIN)\n",
        "valid_dataset = nlp.load_dataset('squad', split=nlp.Split.VALIDATION)\n",
        "\n",
        "# map add_eos_to_examples function to the dataset example wise \n",
        "train_dataset = train_dataset.map(add_eos_to_examples)\n",
        "# map convert_to_features batch wise\n",
        "train_dataset = train_dataset.map(convert_to_features, batched=True)\n",
        "\n",
        "valid_dataset = valid_dataset.map(add_eos_to_examples, load_from_cache_file=False)\n",
        "valid_dataset = valid_dataset.map(convert_to_features, batched=True, load_from_cache_file=False)\n",
        "\n",
        "\n",
        "# set the tensor type and the columns which the dataset should return\n",
        "columns = ['input_ids', 'target_ids', 'attention_mask', 'target_attention_mask']\n",
        "train_dataset.set_format(type='torch', columns=columns)\n",
        "valid_dataset.set_format(type='torch', columns=columns)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "d7fa301d6f7d498bac466849f04321c5",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=4997.0, style=ProgressStyle(description…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "d530b763a0e9415f9b8b1f1787a66619",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=2240.0, style=ProgressStyle(description…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Downloading and preparing dataset squad/plain_text (download: 33.51 MiB, generated: 85.75 MiB, total: 119.27 MiB) to /root/.cache/huggingface/datasets/squad/plain_text/1.0.0...\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "6af39625f78e4e5989739eff7e8849f6",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=8116577.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "741a693169b54851a4cb47369dd9bd1e",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1054280.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "004faf88da144ad2bd3e50b5b9766627",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "3ec93bdea33b40ba82e24b4a0e605bc0",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\rDataset squad downloaded and prepared to /root/.cache/huggingface/datasets/squad/plain_text/1.0.0. Subsequent calls will reuse this data.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "87599it [00:04, 19002.06it/s]\n",
            "100%|██████████| 88/88 [00:50<00:00,  1.75it/s]\n",
            "10570it [00:00, 18815.95it/s]\n",
            "100%|██████████| 11/11 [00:06<00:00,  1.77it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vXJ24xVmlMoN",
        "colab_type": "code",
        "outputId": "a12ce45b-09bc-462f-dae4-f2278c9442c5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "len(train_dataset), len(valid_dataset)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(87599, 10570)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UYOvGLdVgoxt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# cach the dataset, so we can load it directly for training\n",
        "\n",
        "torch.save(train_dataset, 'train_data.pt')\n",
        "torch.save(valid_dataset, 'valid_data.pt')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NASBkGrtnbgj",
        "colab_type": "text"
      },
      "source": [
        "For more details on how to use the nlp library check out this [notebook](https://colab.research.google.com/github/huggingface/nlp/blob/master/notebooks/Overview.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tty8vuMBqI5L",
        "colab_type": "text"
      },
      "source": [
        "## Write training script"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "djKKcgN1cvAX",
        "colab_type": "text"
      },
      "source": [
        "Using the `Trainer` is pretty straightforward. Here are the 4 basic steps which are needed to use trainer.\n",
        "\n",
        "1. **Parse the arguments needed**. These are divided in 3 parts for clarity and seperation (TrainingArguments, ModelArguments and DataTrainingArguments).\n",
        "\n",
        "  1. **TrainingArguments**: These are basicaly the training hyperparameters such as learning rate, batch size, weight decay, gradient accumulation steps etc. See all possible arguments [here](https://github.com/huggingface/transformers/blob/master/src/transformers/training_args.py). These are used by the Trainer.\n",
        "\n",
        "  2. **ModelArguments**: These are the arguments for the model that you want to use such as the model_name_or_path, tokenizer_name etc. You'll need these to load the model and tokenizer.\n",
        "\n",
        "  3. **DataTrainingArguments**: These are as the name suggests arguments needed for the dataset. Such as the directory name where your files are stored etc. You'll need these to load/process the dataset.\n",
        "\n",
        "  TrainingArguments are already defined in the `TrainingArguments` class, you'll need to define `ModelArguments` and `DataTrainingArguments` classes for your task.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "2. Load train and eval datasets\n",
        "3. Initialize the `Trainer`\n",
        "\n",
        "    These are the mininum parameters which you'll for initializing `Trainer`. For full list check [here](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py#L107)\n",
        "\n",
        "    ```\n",
        "      model: PreTrainedModel\n",
        "      args: TrainingArguments\n",
        "      train_dataset: Optional[Dataset]\n",
        "      eval_dataset: Optional[Dataset]\n",
        "    ```\n",
        "4. Start training with  `trainer.train`\n",
        "\n",
        "    Call `trainer.train` and let the magic begin!\n",
        "\n",
        "\n",
        "There are lots of things which the trainer handles for you out of the box such as gradient_accumulation, fp16 training, setting up the optimizer and scheduler, logging with wandb etc. I didn't set-up wandb for this experiment, but will explore it for sure in future experiment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KdmKlMkfcLa0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import dataclasses\n",
        "import logging\n",
        "import os\n",
        "import sys\n",
        "from dataclasses import dataclass, field\n",
        "from typing import Dict, List, Optional\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "from transformers import T5ForConditionalGeneration, T5Tokenizer, EvalPrediction\n",
        "from transformers import (\n",
        "    HfArgumentParser,\n",
        "    DataCollator,\n",
        "    Trainer,\n",
        "    TrainingArguments,\n",
        "    set_seed,\n",
        ")\n",
        "\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "# prepares lm_labels from target_ids, returns examples with keys as expected by the forward method\n",
        "# this is necessacry because the trainer directly passes this dict as arguments to the model\n",
        "# so make sure the keys match the parameter names of the forward method\n",
        "@dataclass\n",
        "class T2TDataCollator(DataCollator):\n",
        "    def collate_batch(self, batch: List) -> Dict[str, torch.Tensor]:\n",
        "        \"\"\"\n",
        "        Take a list of samples from a Dataset and collate them into a batch.\n",
        "        Returns:\n",
        "            A dictionary of tensors\n",
        "        \"\"\"\n",
        "        input_ids = torch.stack([example['input_ids'] for example in batch])\n",
        "        lm_labels = torch.stack([example['target_ids'] for example in batch])\n",
        "        lm_labels[lm_labels[:, :] == 0] = -100\n",
        "        attention_mask = torch.stack([example['attention_mask'] for example in batch])\n",
        "        decoder_attention_mask = torch.stack([example['target_attention_mask'] for example in batch])\n",
        "        \n",
        "\n",
        "        return {\n",
        "            'input_ids': input_ids, \n",
        "            'attention_mask': attention_mask,\n",
        "            'lm_labels': lm_labels, \n",
        "            'decoder_attention_mask': decoder_attention_mask\n",
        "        }\n",
        "\n",
        "\n",
        "@dataclass\n",
        "class ModelArguments:\n",
        "    \"\"\"\n",
        "    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.\n",
        "    \"\"\"\n",
        "\n",
        "    model_name_or_path: str = field(\n",
        "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"}\n",
        "    )\n",
        "    tokenizer_name: Optional[str] = field(\n",
        "        default=None, metadata={\"help\": \"Pretrained tokenizer name or path if not the same as model_name\"}\n",
        "    )\n",
        "    cache_dir: Optional[str] = field(\n",
        "        default=None, metadata={\"help\": \"Where do you want to store the pretrained models downloaded from s3\"}\n",
        "    )\n",
        "\n",
        "@dataclass\n",
        "class DataTrainingArguments:\n",
        "    \"\"\"\n",
        "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
        "    \"\"\"\n",
        "    train_file_path: Optional[str] = field(\n",
        "        default='train_data.pt',\n",
        "        metadata={\"help\": \"Path for cached train dataset\"},\n",
        "    )\n",
        "    valid_file_path: Optional[str] = field(\n",
        "        default='valid_data.pt',\n",
        "        metadata={\"help\": \"Path for cached valid dataset\"},\n",
        "    )\n",
        "    max_len: Optional[int] = field(\n",
        "        default=512,\n",
        "        metadata={\"help\": \"Max input length for the source text\"},\n",
        "    )\n",
        "    target_max_len: Optional[int] = field(\n",
        "        default=32,\n",
        "        metadata={\"help\": \"Max input length for the target text\"},\n",
        "    )\n",
        "\n",
        "\n",
        "def main():\n",
        "    # See all possible arguments in src/transformers/training_args.py\n",
        "    # or by passing the --help flag to this script.\n",
        "    # We now keep distinct sets of args, for a cleaner separation of concerns.\n",
        "\n",
        "    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))\n",
        "\n",
        "    # we will load the arguments from a json file, \n",
        "    #make sure you save the arguments in at ./args.json\n",
        "    model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath('args.json'))\n",
        "\n",
        "    if (\n",
        "        os.path.exists(training_args.output_dir)\n",
        "        and os.listdir(training_args.output_dir)\n",
        "        and training_args.do_train\n",
        "        and not training_args.overwrite_output_dir\n",
        "    ):\n",
        "        raise ValueError(\n",
        "            f\"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome.\"\n",
        "        )\n",
        "\n",
        "    # Setup logging\n",
        "    logging.basicConfig(\n",
        "        format=\"%(asctime)s - %(levelname)s - %(name)s -   %(message)s\",\n",
        "        datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
        "        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,\n",
        "    )\n",
        "    logger.warning(\n",
        "        \"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s\",\n",
        "        training_args.local_rank,\n",
        "        training_args.device,\n",
        "        training_args.n_gpu,\n",
        "        bool(training_args.local_rank != -1),\n",
        "        training_args.fp16,\n",
        "    )\n",
        "    logger.info(\"Training/evaluation parameters %s\", training_args)\n",
        "\n",
        "    # Set seed\n",
        "    set_seed(training_args.seed)\n",
        "\n",
        "    # Load pretrained model and tokenizer\n",
        "    #\n",
        "    # Distributed training:\n",
        "    # The .from_pretrained methods guarantee that only one local process can concurrently\n",
        "    # download model & vocab.\n",
        "\n",
        "    tokenizer = T5Tokenizer.from_pretrained(\n",
        "        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,\n",
        "        cache_dir=model_args.cache_dir,\n",
        "    )\n",
        "    model = T5ForConditionalGeneration.from_pretrained(\n",
        "        model_args.model_name_or_path,\n",
        "        cache_dir=model_args.cache_dir,\n",
        "    )\n",
        "\n",
        "    # Get datasets\n",
        "    print('loading data')\n",
        "    train_dataset  = torch.load(data_args.train_file_path)\n",
        "    valid_dataset = torch.load(data_args.valid_file_path)\n",
        "    print('loading done')\n",
        "\n",
        "    # Initialize our Trainer\n",
        "    trainer = Trainer(\n",
        "        model=model,\n",
        "        args=training_args,\n",
        "        train_dataset=train_dataset,\n",
        "        eval_dataset=valid_dataset,\n",
        "        data_collator=T2TDataCollator(),\n",
        "        prediction_loss_only=True\n",
        "    )\n",
        "\n",
        "    # Training\n",
        "    if training_args.do_train:\n",
        "        trainer.train(\n",
        "            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None\n",
        "        )\n",
        "        trainer.save_model()\n",
        "        # For convenience, we also re-save the tokenizer to the same directory,\n",
        "        # so that you can share your model easily on huggingface.co/models =)\n",
        "        if trainer.is_world_master():\n",
        "            tokenizer.save_pretrained(training_args.output_dir)\n",
        "\n",
        "    # Evaluation\n",
        "    results = {}\n",
        "    if training_args.do_eval and training_args.local_rank in [-1, 0]:\n",
        "        logger.info(\"*** Evaluate ***\")\n",
        "\n",
        "        eval_output = trainer.evaluate()\n",
        "\n",
        "        output_eval_file = os.path.join(training_args.output_dir, \"eval_results.txt\")\n",
        "        with open(output_eval_file, \"w\") as writer:\n",
        "            logger.info(\"***** Eval results *****\")\n",
        "            for key in sorted(eval_output.keys()):\n",
        "                logger.info(\"  %s = %s\", key, str(eval_output[key]))\n",
        "                writer.write(\"%s = %s\\n\" % (key, str(eval_output[key])))\n",
        "    \n",
        "        results.update(eval_output)\n",
        "    \n",
        "    return results\n",
        "\n",
        "\n",
        "def _mp_fn(index):\n",
        "    # For xla_spawn (TPUs)\n",
        "    main()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15duw24hqMBy",
        "colab_type": "text"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n1I6IhBM1KV2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import json"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zOvs9RUllLTw",
        "colab_type": "text"
      },
      "source": [
        "Let's write the arguments in a dict and store in a json file. The above code will load this file and parse the arguments."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ObtXlBVuJqv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "args_dict = {\n",
        "  \"num_cores\": 8,\n",
        "  'training_script': 'train_t5_squad.py',\n",
        "  \"model_name_or_path\": 't5-base',\n",
        "  \"max_len\": 512 ,\n",
        "  \"target_max_len\": 16,\n",
        "  \"output_dir\": './models/tpu',\n",
        "  \"overwrite_output_dir\": True,\n",
        "  \"per_gpu_train_batch_size\": 8,\n",
        "  \"per_gpu_eval_batch_size\": 8,\n",
        "  \"gradient_accumulation_steps\": 4,\n",
        "  \"learning_rate\": 1e-4,\n",
        "  \"tpu_num_cores\": 8,\n",
        "  \"num_train_epochs\": 4,\n",
        "  \"do_train\": True\n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xU5MI8ju1L3w",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with open('args.json', 'w') as f:\n",
        "  json.dump(args_dict, f)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AsQB1Kpjlltp",
        "colab_type": "text"
      },
      "source": [
        "Start training!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UnGuDVPYuyo4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch_xla.distributed.xla_multiprocessing as xmp"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X9_Go99fvW-z",
        "colab_type": "code",
        "outputId": "784c2b84-ecb2-4a2a-871b-dc63f63ccc74",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "1d197b0946534a91b2eb4595baade174",
            "55c18e9169ab44c981de4cb2ebaea578",
            "21065bd4f9fa4d97a20a25c9a35b7737",
            "88d28cee2f02428aa3c4c22a873d6385",
            "12d2d5b9a3f54b25aa8fef6439625db6",
            "32f70451efa6473a8a00807dd22a3248",
            "17c358e37fc346a9ac8ce475599e495a",
            "947afd4c522443b38bc223114178c7ee",
            "a558cfceca8b43568dfad2267a534a35",
            "7a5270ddb66447af9fba298e830fe8a4",
            "d961f9fa40a04ceda3e1962eef4d33cc",
            "d9363dcd670643e6ad2d0b65c4238b98",
            "b25b3d9394be4ac786678181c800e926",
            "3fed14bdec8b4162887d77dea7a9f289",
            "5051978cf84f4f8191c1e3ad4661776e",
            "f67b49a0d48f4714a124bc972e9afa0f",
            "2c4d7bdf7f6c48f0825ef87a12399f39",
            "391fc0937c3b48f8b4c3b33a1c3849a0",
            "11c6ff7dfb19436c920c5114734c0ec9",
            "3b3722a55e3d43bc93c29a505895a25c",
            "f46743ac2da84d6ca77180d8686ead6b",
            "1b0e2b414d6843ae9c918db417ffff4d",
            "e7559a6907f7479db7764f80c4501993",
            "5f89a141b34d478e861ea6bd4a17dcd2",
            "a9a7aeae6fa74b67a349a69e845f8898",
            "0a9d6c739f984b91a0e17724707eb7e1",
            "1ed37d3716374294b072524937625965",
            "4b44de48aeeb4b8dbc9388fc7f7da227",
            "051052ef94d447a4b0233225d8f5ae81",
            "3e43d0668e6f461c9cc461aa735a4437",
            "43e265f386ca4dfd88e226a87f32775c",
            "6568a6e42f734c29ad5c26f2397a5a69",
            "c227333162114432a38b1480ec837701",
            "71a3e5b846a74f409dfd029ea3bf5d61",
            "8a27b38cd17649b1bc720a33050dc520",
            "16497f2b4de443bb9bb14ccaf0c94803",
            "f93601a0c6d84d8b9324fa2380e14d25",
            "f30f39adee694c55bcbaa850cf031b69",
            "ea4e4964058945f287eb2a6271b5bae7",
            "1e74937c1f304d15ad49230b2428590d"
          ]
        }
      },
      "source": [
        "xmp.spawn(_mp_fn, args=(), nprocs=8, start_method='fork')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:27 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:27 - WARNING - __main__ -   Process rank: -1, device: xla:1, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:27 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:27 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:27 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:27 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:27 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:29 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:29 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:29 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:29 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:29 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:29 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:29 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:29 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:29 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:29 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:29 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:29 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:29 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:29 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:29 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:29 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:29 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:29 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:29 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:29 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:29 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:29 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:30 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:30 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:30 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:30 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:30 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:30 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:30 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:30 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:30 - INFO - transformers.training_args -   PyTorch: setting up devices\n",
            "05/16/2020 09:42:30 - WARNING - __main__ -   Process rank: -1, device: xla:0, n_gpu: 0, distributed training: False, 16-bits training: False\n",
            "05/16/2020 09:42:30 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='./models/tpu', overwrite_output_dir=True, do_train=True, do_eval=False, do_predict=False, evaluate_during_training=False, per_gpu_train_batch_size=8, per_gpu_eval_batch_size=8, gradient_accumulation_steps=4, learning_rate=0.0001, weight_decay=0.0, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=4, max_steps=-1, warmup_steps=0, logging_dir=None, logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=8, tpu_metrics_debug=False)\n",
            "05/16/2020 09:42:30 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:30 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:30 - INFO - transformers.tokenization_utils -   loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "05/16/2020 09:42:30 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "\n",
            "05/16/2020 09:42:30 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/t5-base-config.json from cache at /root/.cache/torch/transformers/40578967d1f029acb6162b36db9d8b4307063e885990ccd297c2c5be1cf1b3d7.2995d650f5eba18c8baa4146e210d32d56165e90d374281741fc78b872cd6c9b\n",
            "05/16/2020 09:42:30 - INFO - transformers.configuration_utils -   Model config T5Config {\n",
            "  \"architectures\": [\n",
            "    \"T5WithLMHeadModel\"\n",
            "  ],\n",
            "  \"d_ff\": 3072,\n",
            "  \"d_kv\": 64,\n",
            "  \"d_model\": 768,\n",
            "  \"decoder_start_token_id\": 0,\n",
            "  \"dropout_rate\": 0.1,\n",
            "  \"eos_token_id\": 1,\n",
            "  \"initializer_factor\": 1.0,\n",
            "  \"is_encoder_decoder\": true,\n",
            "  \"layer_norm_epsilon\": 1e-06,\n",
            "  \"model_type\": \"t5\",\n",
            "  \"n_positions\": 512,\n",
            "  \"num_heads\": 12,\n",
            "  \"num_layers\": 12,\n",
            "  \"output_past\": true,\n",
            "  \"pad_token_id\": 0,\n",
            "  \"relative_attention_num_buckets\": 32,\n",
            "  \"task_specific_params\": {\n",
            "    \"summarization\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"length_penalty\": 2.0,\n",
            "      \"max_length\": 200,\n",
            "      \"min_length\": 30,\n",
            "      \"no_repeat_ngram_size\": 3,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"summarize: \"\n",
            "    },\n",
            "    \"translation_en_to_de\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to German: \"\n",
            "    },\n",
            "    \"translation_en_to_fr\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to French: \"\n",
            "    },\n",
            "    \"translation_en_to_ro\": {\n",
            "      \"early_stopping\": true,\n",
            "      \"max_length\": 300,\n",
            "      \"num_beams\": 4,\n",
            "      \"prefix\": \"translate English to Romanian: \"\n",
            "    }\n",
            "  },\n",
            "  \"vocab_size\": 32128\n",
            "}\n",
            "\n",
            "05/16/2020 09:42:30 - INFO - transformers.modeling_utils -   loading weights file https://cdn.huggingface.co/t5-base-pytorch_model.bin from cache at /root/.cache/torch/transformers/f6f2fde9fa7611f4eff74620de9cbe734e7a717b5b143bd283cae4c2d6022990.54f906ff53bd09195cfc183a29cadc81b7705f07fcdb796d24163cb632b6bdfa\n",
            "05/16/2020 09:42:34 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:34 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:34 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n",
            "05/16/2020 09:42:37 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n",
            "05/16/2020 09:42:37 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n",
            "05/16/2020 09:42:37 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n",
            "05/16/2020 09:42:37 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:37 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n",
            "05/16/2020 09:42:37 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:38 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:38 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:38 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n",
            "05/16/2020 09:42:38 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:38 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:38 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n",
            "05/16/2020 09:42:38 - INFO - transformers.modeling_utils -   Weights of T5ForConditionalGeneration not initialized from pretrained model: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight', 'lm_head.weight']\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading data\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:38 - INFO - nlp.utils.file_utils -   PyTorch version 1.6.0a0+bf2bbd9 available.\n",
            "05/16/2020 09:42:38 - INFO - nlp.utils.file_utils -   TensorFlow version 2.2.0 available.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:40 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:40 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:40 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:41 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:41 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:41 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "loading done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:42:41 - INFO - transformers.trainer -   You are instantiating a Trainer but W&B is not installed. To use wandb logging, run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface.\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:20 - INFO - transformers.trainer -     Total optimization steps = 1368\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1d197b0946534a91b2eb4595baade174",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Epoch', max=4.0, style=ProgressStyle(description_width='i…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a558cfceca8b43568dfad2267a534a35",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Iteration', max=1369.0, style=ProgressStyle(description_w…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 09:43:33 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:33 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:33 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:33 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:33 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:33 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:33 - INFO - transformers.trainer -     Total optimization steps = 1368\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:36 - INFO - transformers.trainer -     Total optimization steps = 1368\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:40 - INFO - transformers.trainer -     Total optimization steps = 1368\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:44 - INFO - transformers.trainer -     Total optimization steps = 1368\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:45 - INFO - transformers.trainer -     Total optimization steps = 1368\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:46 - INFO - transformers.trainer -     Total optimization steps = 1368\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -   ***** Running training *****\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -     Num examples = 87599\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -     Num Epochs = 4\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -     Instantaneous batch size per device = 8\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -     Total train batch size (w. parallel, distributed & accumulation) = 64\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -     Gradient Accumulation steps = 4\n",
            "05/16/2020 09:43:47 - INFO - transformers.trainer -     Total optimization steps = 1368\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "2c4d7bdf7f6c48f0825ef87a12399f39",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Iteration', max=1369.0, style=ProgressStyle(description_w…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n",
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n",
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n",
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a9a7aeae6fa74b67a349a69e845f8898",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Iteration', max=1369.0, style=ProgressStyle(description_w…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "c227333162114432a38b1480ec837701",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Iteration', max=1369.0, style=ProgressStyle(description_w…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n",
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n",
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:44 - INFO - transformers.trainer -   \n",
            "\n",
            "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
            "\n",
            "\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.trainer -   Saving model checkpoint to ./models/tpu\n",
            "05/16/2020 10:03:45 - INFO - transformers.configuration_utils -   Configuration saved in ./models/tpu/config.json\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:54 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n",
            "05/16/2020 10:03:55 - INFO - transformers.modeling_utils -   Model weights saved in ./models/tpu/pytorch_model.bin\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LAe5vnbi-dyx",
        "colab_type": "text"
      },
      "source": [
        "## Eval"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VxWscyGVl05C",
        "colab_type": "text"
      },
      "source": [
        "There are two gotchas here. First the metrics functionality in the nlp package is still work-in-progress so we will use the official squad evaluation script. Second, for some reason which I couldn't figure out, the `.generate` method is not working on TPU so will need to do prediction on CPU. For predicting the validation set it almost takes 40 mins."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H8AbD1B7TR0k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## SQuAD evaluation script. Modifed slightly for this notebook\n",
        "\n",
        "from __future__ import print_function\n",
        "from collections import Counter\n",
        "import string\n",
        "import re\n",
        "import argparse\n",
        "import json\n",
        "import sys\n",
        "\n",
        "\n",
        "def normalize_answer(s):\n",
        "    \"\"\"Lower text and remove punctuation, articles and extra whitespace.\"\"\"\n",
        "    def remove_articles(text):\n",
        "        return re.sub(r'\\b(a|an|the)\\b', ' ', text)\n",
        "\n",
        "    def white_space_fix(text):\n",
        "        return ' '.join(text.split())\n",
        "\n",
        "    def remove_punc(text):\n",
        "        exclude = set(string.punctuation)\n",
        "        return ''.join(ch for ch in text if ch not in exclude)\n",
        "\n",
        "    def lower(text):\n",
        "        return text.lower()\n",
        "\n",
        "    return white_space_fix(remove_articles(remove_punc(lower(s))))\n",
        "\n",
        "\n",
        "def f1_score(prediction, ground_truth):\n",
        "    prediction_tokens = normalize_answer(prediction).split()\n",
        "    ground_truth_tokens = normalize_answer(ground_truth).split()\n",
        "    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n",
        "    num_same = sum(common.values())\n",
        "    if num_same == 0:\n",
        "        return 0\n",
        "    precision = 1.0 * num_same / len(prediction_tokens)\n",
        "    recall = 1.0 * num_same / len(ground_truth_tokens)\n",
        "    f1 = (2 * precision * recall) / (precision + recall)\n",
        "    return f1\n",
        "\n",
        "\n",
        "def exact_match_score(prediction, ground_truth):\n",
        "    return (normalize_answer(prediction) == normalize_answer(ground_truth))\n",
        "\n",
        "\n",
        "def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):\n",
        "    scores_for_ground_truths = []\n",
        "    for ground_truth in ground_truths:\n",
        "        score = metric_fn(prediction, ground_truth)\n",
        "        scores_for_ground_truths.append(score)\n",
        "    return max(scores_for_ground_truths)\n",
        "\n",
        "\n",
        "def evaluate(gold_answers, predictions):\n",
        "    f1 = exact_match = total = 0\n",
        "\n",
        "    for ground_truths, prediction in zip(gold_answers, predictions):\n",
        "      total += 1\n",
        "      exact_match += metric_max_over_ground_truths(\n",
        "                    exact_match_score, prediction, ground_truths)\n",
        "      f1 += metric_max_over_ground_truths(\n",
        "          f1_score, prediction, ground_truths)\n",
        "    \n",
        "    exact_match = 100.0 * exact_match / total\n",
        "    f1 = 100.0 * f1 / total\n",
        "\n",
        "    return {'exact_match': exact_match, 'f1': f1}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FCZVPIUK9fyn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch_xla\n",
        "import torch_xla.core.xla_model as xm\n",
        "\n",
        "import nlp\n",
        "from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
        "\n",
        "from tqdm.auto import tqdm"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vmV9c-w39c7C",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = T5ForConditionalGeneration.from_pretrained('models/tpu').to('cpu') # because its loaded on xla by default\n",
        "tokenizer = T5Tokenizer.from_pretrained('models/tpu')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vjstID24-IAw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "valid_dataset = torch.load('valid_data.pt')\n",
        "dataloader = torch.utils.data.DataLoader(valid_dataset, batch_size=32)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v7TUzb-T-YtF",
        "colab_type": "code",
        "outputId": "c4d045a2-fd25-476d-9461-68571523b861",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "df609e7f302e48949a49fae1b0deeb54",
            "be907685a2b949d0a088a907a2cb2a90",
            "743bba68cbe943d78ce63a1efd70d929",
            "f3fffc31462b49e193b3c24d600b4ddb",
            "87fa58101b24400da90e4b0f87c7cb1a",
            "3c39d97054ba47a3ad3d9bb9a882756e",
            "20855b3cb56944cfb0d6aa13bd06429a",
            "779d80cc14104171ab40ff53170e4807"
          ]
        }
      },
      "source": [
        "answers = []\n",
        "for batch in tqdm(dataloader):\n",
        "  outs = model.generate(input_ids=batch['input_ids'], \n",
        "                        attention_mask=batch['attention_mask'],\n",
        "                        max_length=16,\n",
        "                        early_stopping=True)\n",
        "  outs = [tokenizer.decode(ids) for ids in outs]\n",
        "  answers.extend(outs)"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "df609e7f302e48949a49fae1b0deeb54",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=331.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VJ7CFrEtLD4F",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "predictions = []\n",
        "references = []\n",
        "for ref, pred in zip(valid_dataset, answers):\n",
        "  predictions.append(pred)\n",
        "  references.append(ref['answers']['text'])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BDpnF6NmMWEl",
        "colab_type": "code",
        "outputId": "79f82ae2-8536-4e7c-f0c4-8a1fbcab8aa7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "predictions[0], references[0]"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('Denver Broncos', ['Denver Broncos', 'Denver Broncos', 'Denver Broncos'])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dtRfqm3Odgh1",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "7d6b5e09-63a6-415f-fb95-9485bb8f7fbc"
      },
      "source": [
        "evaluate(references, predictions)"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'exact_match': 81.56102175969725, 'f1': 89.96016967193422}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 27
        }
      ]
    }
  ]
}